import cv2
import torch
from typing import List
import os

count = 0


def print_bbox_in_img(img_meta, gt_bboxes, bboxes, max_overlaps, saved_path):
    """_print_bbox_in_img.

    Args:
        img_meta: img_meta infomation
        gt_bboxes: gt bboxes location
        bboxes: all generated anchors
        max_overlaps: a flag that indicates a bbox being assigned as pos anchor.
    """

    INF = 100000000
    img_raw = cv2.imread(img_meta["filename"], cv2.IMREAD_COLOR)

    global count
    #  print(count)
    count = count + 1
    if count > 100:
        os._exit(0)

    for i in range(len(gt_bboxes)):
        b = gt_bboxes[i, :]
        # print(b)
        cv2.rectangle(
            img_raw, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])), (0, 0, 255), 2
        )
    bboxes_assign = bboxes[(max_overlaps != -INF), :]
    for i in range(len(bboxes_assign)):
        b = bboxes_assign[i, :]
        # print(b)
        cv2.rectangle(
            img_raw, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])), (0, 255, 0), 2
        )
    cv2.imwrite(os.path.join(saved_path, str(count) + ".jpg"), img_raw)


def print_num_anchor(
    bboxes,
    gt_bboxes,
    num_bboxes,
    num_level_bboxes,
    index,
    img_meta,
    gt_labels,
    saved_path,
    overlaps_thr_per_gt,
):
    """_print_num_anchor.

    Args:
        bboxes: anchors
        gt_bboxes: gt bboxes location
        num_bboxes: the amount of generated anchors
        num_level_bboxes(List): the amount of generated anchors from each FPN output
        index: the positive anchor indexes
        img_meta: img_meta infomation
        gt_labels: gt labels
        overlaps_thr_per_gt: threhold for positive sample in every gt bbox.
    """
    num_gt = gt_bboxes.size(0)
    num_bboxes_from_level, pos_bboxes = _get_num_bboxes_from_level(
        bboxes, num_gt, num_bboxes, num_level_bboxes, index
    )
    str_pos_bboxes = [[i.__str__() for i in pos_bbox] for pos_bbox in pos_bboxes]
    thr_per_gt_list = overlaps_thr_per_gt.tolist()
    msg_type = "full"

    if msg_type == "full":
        written_msgs = ""
        written_msgs += f"{img_meta['filename']}\n"
        written_msgs += f"\tTotally {num_gt} GT Objects\n"
        written_msgs += (
            f"\tThreshold for gts:{[round(i, 3) for i in thr_per_gt_list]}\n"
        )
        for i in range(num_gt):
            written_msgs += f"\t\t{i}th GT cls: {int(gt_labels[i].cpu())}. "
            written_msgs += "num pos anchor in fpn level"
            written_msgs += f"{num_bboxes_from_level[i]}. "
            area = (gt_bboxes[i][2] - gt_bboxes[i][0]) * (
                gt_bboxes[i][3] - gt_bboxes[i][1]
            )
            area = float(area.cpu())
            written_msgs += f"area: {area:.2f}\n"

            # print pos anchor assigned to i-th gt
            written_msgs += "\t\t\t"
            written_msgs += "\n\t\t\t".join(str_pos_bboxes[i])
            written_msgs += "\n"
    elif msg_type == "only_look_level1":
        written_msgs = ""
        written_msgs += f"{img_meta['filename']}\n"
        written_msgs += f"\tTotally {num_gt} GT Objects\n"
        written_msgs += f"Threshold for gts:{thr_per_gt_list}\n"
        for i in range(num_gt):
            written_msgs += f"\t\t{i}th GT cls: {int(gt_labels[i].cpu())}. "
            written_msgs += "num pos anchor in 1st level: "
            written_msgs += f"{num_bboxes_from_level[i][0]}. "
            area = (gt_bboxes[i][2] - gt_bboxes[i][0]) * (
                gt_bboxes[i][3] - gt_bboxes[i][1]
            )
            area = float(area.cpu())
            written_msgs += f"area: {area:.2f}\n"

    #
    saved_path = saved_path.split("/")
    filename_index = saved_path[-1]
    saved_path = "/".join(saved_path[:-1])
    with open(os.path.join(saved_path, f"data{filename_index}.txt"), "a") as f:
        f.write(written_msgs)


def _get_num_bboxes_from_level(bboxes, num_gt, num_bboxes, num_level_bboxes, index):
    """_get_num_bboxes_from_level.

    Args:
        bboxes: anchors, predefined bboxes
        num_gt: the amount of Groundtruth bboxes
        num_bboxes: the amount of bboxes
        num_level_bboxes(List): the i-th element denotes the amount of anchors in
            i-th FPN featmap
        index: index of positive anchors

    Returns:
        num_bboxes_from_level (List(List)): the amout of positive anchors from all
            fpn level to every gt bboxes. num_bboxes_from_level[i][j] denotes gt[i]
            in FPN-level[j] has this amount of anchors.
        pos_bboxes (List(List)): the positive bboxes locations s.t gt-bboxes and FPN
            level. pos_bboxes[i][j] denotes all pos bboxes that gt[i] has in
            FPN-level[j].
    """
    num_bboxes_from_level = [[0 for _ in num_level_bboxes] for _ in range(num_gt)]
    pos_bboxes = [[] for _ in range(num_gt)]
    for i in index:
        gt_index = i // num_bboxes
        i = i % num_bboxes
        pos_bboxes[gt_index].append(bboxes[i])

        for which_level, num_level_bbox in enumerate(num_level_bboxes):
            i -= num_level_bbox
            if i < 0:
                num_bboxes_from_level[gt_index][which_level] += 1
                break

    return (
        [num_bboxes_from_level[i][:] for i in range(len(num_bboxes_from_level))],
        pos_bboxes,
    )


if __name__ == "__main__":
    import torch

    pos_bboxes = [[]]
    pos_bboxes[0].append(torch.ones(3, 3))
    pos_bboxes[0].append(torch.ones(4, 4))
    pos_bboxes[0].append(torch.ones(3, 3))
    print(pos_bboxes)
